package com.nidu.demo.permission.handler;

import cn.hutool.core.collection.CollUtil;
import com.baomidou.mybatisplus.extension.plugins.handler.DataPermissionHandler;
import com.nidu.demo.common.enums.DataScopeEnum;
import com.nidu.demo.permission.config.DataPermissionContextHolder;
import com.nidu.demo.permission.config.DataPermissionProperties;
import com.nidu.demo.security.util.SecurityUtils;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.LongValue;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
import net.sf.jsqlparser.expression.operators.relational.InExpression;
import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
import net.sf.jsqlparser.schema.Column;
import lombok.extern.slf4j.Slf4j;

import java.util.HashSet;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * MyBatis数据权限处理器
 * 参考租户系统的MybatisTenantLineHandler实现
 */
@Slf4j
public class MybatisDataPermissionHandler implements DataPermissionHandler {

    private final DataPermissionProperties properties;
    private final Set<String> ignoreTables = new HashSet<>();

    public MybatisDataPermissionHandler(DataPermissionProperties properties) {
        this.properties = properties;
        // 不同 DB 下，大小写的习惯不同，所以需要都添加进去
        properties.getIgnoreTables().forEach(table -> {
            ignoreTables.add(table.toLowerCase());
            ignoreTables.add(table.toUpperCase());
        });
    }

    @Override
    public Expression getSqlSegment(Expression where, String mappedStatementId) {
        // 如果忽略数据权限或没有设置数据权限上下文，直接返回原条件
        if (DataPermissionContextHolder.isIgnore()) {
            return where;
        }

        DataScopeEnum dataScope = DataPermissionContextHolder.getDataScope();
        if (dataScope == null) {
            return where;
        }

        Expression dataPermissionExpression = buildDataPermissionExpression(dataScope);
        if (dataPermissionExpression == null) {
            return where;
        }

        // 将数据权限条件与原有条件组合
        if (where == null) {
            return dataPermissionExpression;
        } else {
            return new AndExpression(where, dataPermissionExpression);
        }
    }

    /**
     * 构建数据权限表达式
     */
    private Expression buildDataPermissionExpression(DataScopeEnum dataScope) {
        switch (dataScope) {
            case SELF:
                // 仅本人数据权限，添加用户ID限制
                return buildUserPermissionExpression();
            case DEPT_CUSTOM:
            case DEPT_ONLY:
            case DEPT_AND_CHILD:
                // 部门相关权限，添加部门ID限制
                return buildDeptPermissionExpression();
            case ALL:
                // 全部数据权限，不添加限制
                return null;
            default:
                log.warn("未知的数据权限范围: {}", dataScope);
                return null;
        }
    }

    /**
     * 构建部门权限表达式
     */
    private Expression buildDeptPermissionExpression() {
        Set<Long> deptIds = DataPermissionContextHolder.getDeptIds();
        if (CollUtil.isEmpty(deptIds)) {
            return null;
        }

        Column deptColumn = new Column(properties.getDeptColumnName());
        
        if (deptIds.size() == 1) {
            // 单个部门，使用等号
            EqualsTo equalsTo = new EqualsTo();
            equalsTo.setLeftExpression(deptColumn);
            equalsTo.setRightExpression(new LongValue(deptIds.iterator().next()));
            return equalsTo;
        } else {
            // 多个部门，使用IN
            InExpression inExpression = new InExpression();
            inExpression.setLeftExpression(deptColumn);
            
            ExpressionList expressionList = new ExpressionList();
            expressionList.setExpressions(
                deptIds.stream()
                    .map(LongValue::new)
                    .collect(Collectors.toList())
            );
            inExpression.setRightExpression(expressionList);
            return inExpression;
        }
    }

    /**
     * 构建用户权限表达式
     */
    private Expression buildUserPermissionExpression() {
        Long userId = SecurityUtils.getLoginUserId();
        if (userId == null) {
            return null;
        }

        Column userColumn = new Column(properties.getUserColumnName());
        EqualsTo equalsTo = new EqualsTo();
        equalsTo.setLeftExpression(userColumn);
        equalsTo.setRightExpression(new LongValue(userId));
        return equalsTo;
    }
}